import torch
import torch.nn as nn
from torch.optim.lr_scheduler import ReduceLROnPlateau
import matplotlib.pyplot as plt

def train_model(model, train_loader, criterion, optimizer, device, num_epochs):
    """
    训练模型并记录训练过程中的损失和准确率
    
    参数:
        model: 要训练的模型
        train_loader: 训练数据加载器
        criterion: 损失函数
        optimizer: 优化器
        device: 训练设备(CPU/GPU)
        num_epochs: 训练轮数
        
    返回:
        model: 训练好的模型
        train_losses: 训练损失记录
        train_accuracies: 训练准确率记录
    """
    # 学习率调度器
    scheduler = ReduceLROnPlateau(optimizer, 'min', patience=3, factor=0.1, verbose=True)
    
    model.train()
    train_losses = []
    train_accuracies = []
    
    for epoch in range(num_epochs):
        running_loss = 0.0
        correct = 0
        total = 0
        
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            
            # 计算训练准确率
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        
        # 计算平均损失和准确率
        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100 * correct / total
        
        # 更新学习率
        scheduler.step(epoch_loss)
        
        train_losses.append(epoch_loss)
        train_accuracies.append(epoch_acc)
        
        print(f"Epoch [{epoch + 1}/{num_epochs}], "
              f"Loss: {epoch_loss:.4f}, "
              f"Accuracy: {epoch_acc:.2f}%, "
              f"LR: {optimizer.param_groups[0]['lr']:.6f}")
    
    # 绘制训练曲线
    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Training Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    
    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label='Training Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.legend()
    
    plt.tight_layout()
    plt.savefig('training_curve.png')
    plt.close()
    
    return model, train_losses, train_accuracies

def evaluate_model(model, test_loader, device, classes=None):
    """
    评估模型在测试集上的性能
    
    参数:
        model: 要评估的模型
        test_loader: 测试数据加载器
        device: 评估设备(CPU/GPU)
        classes: (可选)类别名称列表
        
    返回:
        accuracy: 测试准确率
        test_loss: 测试损失
        class_accuracies: 每个类别的准确率
    """
    model.eval()
    correct = 0
    total = 0
    test_loss = 0.0
    class_correct = [0] * 10
    class_total = [0] * 10
    
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = nn.CrossEntropyLoss()(outputs, labels)
            test_loss += loss.item()
            
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            
            # 计算每个类别的准确率
            for i in range(len(labels)):
                label = labels[i]
                class_correct[label] += (predicted[i] == label).item()
                class_total[label] += 1
    
    accuracy = 100 * correct / total
    test_loss /= len(test_loader)
    
    # 计算每个类别的准确率
    class_accuracies = [100 * class_correct[i] / class_total[i] for i in range(10)]
    
    print(f"Test Loss: {test_loss:.4f}, Test Accuracy: {accuracy:.2f}%")
    
    # 只有在提供了classes参数时才打印类别准确率
    if classes is not None:
        print("Class-wise Accuracies:")
        for i in range(10):
            print(f"{classes[i]}: {class_accuracies[i]:.2f}%")
    
    return accuracy, test_loss, class_accuracies